#ifndef NETWORK_H_
#define NETWORK_H_

#include <arpa/inet.h>
#include <errno.h>
#include <stdarg.h>
#include <stdio.h>
#include <string.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <unistd.h>

namespace network {

class NetworkContext {
 public:
  enum Protocol { NETWORK_PROTOCOL_TCP = 0, NETWORK_PROTOCOL_UDP = 1 };

  enum IoProcessType { VernoxSelect };

  // 增强回调类型：增加用户数据指针
  typedef int (*IOCallback)(void* buffer, size_t length, int socketFd,
                            void* userData);

  // 构造函数
  NetworkContext(const char* addr, uint16_t port, Protocol protocol)
      : addr_(addr),
        port_(port),
        protocol_(protocol),
        sockfd_(-1),
        callback_(0),
        user_data_(0) {
    memset(&server_addr_, 0, sizeof(server_addr_));
    server_addr_.sin_family = AF_INET;
    server_addr_.sin_port = htons(port_);
    inet_pton(AF_INET, addr_, &server_addr_.sin_addr);
  }

  // 禁用拷贝
 private:
  NetworkContext(const NetworkContext&);
  NetworkContext& operator=(const NetworkContext&);

 public:
  ~NetworkContext() { CloseSocket(); }

  void SetCallback(IOCallback callback, void* userData = 0) {
    callback_ = callback;
    user_data_ = userData;
  }

  int StartServer(IoProcessType process_type) {
    int type = (protocol_ == NETWORK_PROTOCOL_TCP) ? SOCK_STREAM : SOCK_DGRAM;
    sockfd_ = socket(AF_INET, type, 0);
    if (sockfd_ == -1) {
      LogError("socket() failed: ");
      return -1;
    }

    // 端口复用
    int opt = 1;
    setsockopt(sockfd_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    if (bind(sockfd_, (struct sockaddr*)&server_addr_, sizeof(server_addr_))) {
      LogError("bind() failed: ");
      CloseSocket();
      return -1;
    }

    if (protocol_ == NETWORK_PROTOCOL_TCP) {
      if (listen(sockfd_, 5) == -1) {
        LogError("listen() failed: ");
        CloseSocket();
        return -1;
      }
    }

    LogMessage("Server started on %s:%d", addr_, port_);

    if (process_type == VernoxSelect) {
      RunServer();
    }
    return 0;
  }

  int StartClient() {
    int type = (protocol_ == NETWORK_PROTOCOL_TCP) ? SOCK_STREAM : SOCK_DGRAM;
    sockfd_ = socket(AF_INET, type, 0);
    if (sockfd_ == -1) {
      LogError("socket() failed: ");
      return -1;
    }

    if (protocol_ == NETWORK_PROTOCOL_TCP) {
      if (connect(sockfd_, (struct sockaddr*)&server_addr_,
                  sizeof(server_addr_))) {
        LogError("connect() failed: ");
        CloseSocket();
        return -1;
      }
      LogMessage("TCP Connected to %s:%d", addr_, port_);
    } else {
      LogMessage("UDP Ready to send to %s:%d", addr_, port_);
    }
    return 0;
  }

  int Send(const void* buff, size_t size) {
    if (sockfd_ == -1) return -1;

    if (protocol_ == NETWORK_PROTOCOL_TCP) {
      return send(sockfd_, buff, size, 0);
    } else {
      return sendto(sockfd_, buff, size, 0, (struct sockaddr*)&server_addr_,
                    sizeof(server_addr_));
    }
  }

  int Recv(void* buff, size_t size) {
    if (sockfd_ == -1) return -1;

    if (protocol_ == NETWORK_PROTOCOL_TCP) {
      return recv(sockfd_, buff, size, 0);
    } else {
      struct sockaddr_in from_addr;
      socklen_t addr_len = sizeof(from_addr);
      return recvfrom(sockfd_, buff, size, 0, (struct sockaddr*)&from_addr,
                      &addr_len);
    }
  }

  static int Send(int fd, const void* buff, size_t size) {
    return send(fd, buff, size, 0);
  }

  static int Recv(int fd, void* buff, size_t size) {
    return recv(fd, buff, size, 0);
  }

 private:
  void CloseSocket() {
    if (sockfd_ != -1) {
      close(sockfd_);
      sockfd_ = -1;
    }
  }

  void RunServer() {
    fd_set read_fds;
    int client_fds[FD_SETSIZE] = {0};
    int max_fd = sockfd_;

    while (1) {
      FD_ZERO(&read_fds);
      FD_SET(sockfd_, &read_fds);
      for (int i = 0; i < FD_SETSIZE; ++i) {
        if (client_fds[i] > 0) {
          FD_SET(client_fds[i], &read_fds);
          if (client_fds[i] > max_fd) max_fd = client_fds[i];
        }
      }

      int activity = select(max_fd + 1, &read_fds, 0, 0, 0);
      if (activity < 0) {
        LogError("select() failed: ");
        break;
      }

      // 处理新连接（仅TCP）
      if (protocol_ == NETWORK_PROTOCOL_TCP && FD_ISSET(sockfd_, &read_fds)) {
        struct sockaddr_in client_addr;
        socklen_t addr_len = sizeof(client_addr);
        int new_fd = accept(sockfd_, (struct sockaddr*)&client_addr, &addr_len);
        if (new_fd == -1) {
          LogError("accept() failed: ");
        } else {
          LogMessage("New client fd=%d", new_fd);
          // 添加到客户端列表
          int added = 0;
          for (int i = 0; i < FD_SETSIZE; ++i) {
            if (client_fds[i] == 0) {
              client_fds[i] = new_fd;
              added = 1;
              break;
            }
          }
          if (!added) {
            LogMessage("Too many clients, close %d", new_fd);
            close(new_fd);
          }
        }
      }

      // 处理数据
      for (int i = 0; i < FD_SETSIZE; ++i) {
        int fd = client_fds[i];
        if (fd <= 0) continue;

        if (FD_ISSET(fd, &read_fds)) {
          char buffer[1024];
          int bytes = Recv(fd, buffer, sizeof(buffer));
          if (bytes > 0) {
            if (callback_) {
              callback_(buffer, bytes, fd, user_data_);
            }
          } else {
            if (bytes == 0) {
              LogMessage("Client %d disconnected", fd);
            } else {
              LogError("recv() error: ");
            }
            close(fd);
            client_fds[i] = 0;
          }
        }
      }
    }
  }

  // 日志函数（可重定向）
  void LogError(const char* prefix) {
    fprintf(stderr, "%s error code: %d\n", prefix, errno);
  }

  void LogMessage(const char* format, ...) {
    va_list args;
    va_start(args, format);
    vprintf(format, args);
    va_end(args);
    printf("\n");
  }

  const char* addr_;
  uint16_t port_;
  Protocol protocol_;
  int sockfd_;
  struct sockaddr_in server_addr_;
  IOCallback callback_;
  void* user_data_;
};

}  // namespace network

#endif  // NETWORK_H_